import math
import random
import sys

import pygame
import pygame.gfxdraw

from utils import rotate_vector2

SCREEN_WIDTH = 1024
SCREEN_HEIGHT = 768

FPS = 60


"""
Cohesion: Each Boid tries to stay close to its neighbors.
Separation: Each Boid tries to maintain a minimum distance from other Boids.
Alignment: Each Boid tries to align itself with the average velocity of its neighbors.
"""

class Boid:
    def __init__(self, boids, game_world):
        self.game_world = game_world
        self.boids = boids
        self.speed = 260
        self.position = pygame.Vector2(random.randint(0, game_world.width), random.randint(0, game_world.height))
        self.velocity = pygame.Vector2(
            random.uniform(-1, 1),
            random.uniform(-1, 1)
        )
        self.perception_radius = 50
        self.r = 5

    def update(self, dt):
        # separation 分离
        sep = self.separation(self.boids)
        # sep = pygame.Vector2(0, 0)
        # alignment 排列
        ali = self.alignment(self.boids)
        # ali = pygame.Vector2(0, 0)
        # cohesion 聚集
        coh = self.cohesion(self.boids)

        # 调整速度
        self.velocity = (self.velocity + sep * 1.5+ ali * 1.5 + coh * 1.5).normalize()

        # 调整位置
        self.position += self.velocity * dt * self.speed

        # 边界处理
        self.cross_border(self.game_world.width, self.game_world.height)

    def draw(self, screen):
        shape = self.get_shape()
        pygame.gfxdraw.filled_polygon(screen, shape, (49, 179, 243))

    def get_shape(self, scale=7):
        head_point = pygame.Vector2(1, 0) + self.position
        if self.velocity.length() > 0:
            head_point = self.velocity.normalize()
        left_point = rotate_vector2(head_point, 135)
        right_point = rotate_vector2(head_point, -135)
        result = [
            head_point * scale + self.position,
            left_point * scale + self.position,
            right_point * scale + self.position
        ]
        return result

    def cross_border(self, gw_width, gw_height):
        if self.position.x < -self.r:
            self.position.x = gw_width + self.r
        if self.position.x > gw_width + self.r:
            self.position.x = -self.r

        if self.position.y < -self.r:
            self.position.y = gw_height + self.r
        if self.position.y > gw_height + self.r:
            self.position.y = -self.r

    def separation(self, boids):
        steering = pygame.Vector2(0, 0)
        count = 0
        for boid in boids:
            distance = (self.position - boid.position).length()
            if distance < self.perception_radius:
                diff = self.position - boid.position
                if distance != 0:
                    steering += diff / distance
                count += 1
        if count > 0:
            steering /= count

        if steering.length() > 0:
            steering.normalize_ip()

        return steering

    def alignment(self, boids):
        steering = pygame.Vector2(0, 0)
        count = 0
        for boid in boids:
            distance = (self.position - boid.position).length()
            if distance < self.perception_radius:
                steering += boid.velocity
                # print(steering)
                count += 1
        if count > 0:
            steering /= count

        if steering.length() > 0:
            steering.normalize_ip()

        return steering

    def cohesion(self, boids):
        steering = pygame.Vector2(0, 0)
        count = 0
        center_of_mass = pygame.Vector2(0, 0)
        for boid in boids:
            distance = (self.position - boid.position).length()
            if distance < self.perception_radius:
                center_of_mass += boid.position
                count += 1
        if count > 0:
            center_of_mass /= count

        diff = center_of_mass - self.position
        steering += diff

        if steering.length() > 0:
            steering.normalize_ip()

        return steering



class GameWorld:
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode()
        self.width, self.height = self.screen.get_size()
        self.clock = pygame.time.Clock()

        self.boids = []
        for i in range(160):
            self.boids.append(Boid(self.boids, self))

    def run(self):
        is_running = True
        while is_running:
            # events
            events = pygame.event.get()
            for event in events:
                if event.type == pygame.QUIT:
                    is_running = False
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        is_running = False

            # clock
            dt = self.clock.tick(FPS) / 1000

            # update
            for b in self.boids:
                b.update(dt)

            # draw
            self.screen.fill((0, 40, 120))

            for b in self.boids:
                b.draw(self.screen)
            # pygame.draw.circle(self.screen, (200,200,200,240),
            #                    self.boids[0].position, 100, 1)

            pygame.display.flip()

    def exit(self):
        pygame.quit()
        sys.exit()


def main():
    gw = GameWorld()
    gw.run()
    gw.exit()


if __name__ == '__main__':
    main()
